
import numpy as np
import mindspore as ms
import mindspore.nn as nn


from model import SPPLeNet5, CaltechDataset, gen_dataset


def res_file(res, res_name='./result.txt'):
    with open(res_name, "w") as f:
        for pred in res:
            f.write("{}\n".format(pred))


def val_resnet_train(net_param_path, file_path, train='train_val'):
    dataset_class = CaltechDataset(file_path, train=True)

    network = SPPLeNet5(num_class=256, num_channel=3, num_layers=9)
    # 加载已经保存的用于测试的模型
    param_dict = ms.load_checkpoint(net_param_path)
    # 加载参数到网络中
    ms.load_param_into_net(network, param_dict)

    net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
    model = ms.Model(network, loss_fn=net_loss, metrics={'accuracy'})

    if train == 'train_val':
        dataset = gen_dataset(dataset_class, train='train_val', batch_size=128)
        acc = model.eval(dataset)
        print(acc)

    elif train == 'train':
        dataset = gen_dataset(dataset_class, train='train_val', batch_size=1)
        dataset = dataset.create_dict_iterator()
        flag = 0
        for img_dict in dataset:
            imgs = img_dict['image']
            img_label = img_dict['label']

            output = model.predict(ms.Tensor(imgs))
            predicted = np.argmax(output.asnumpy(), axis=1)
            print("{}: {}, {}".format(flag, predicted.item(), img_label.item()))
            flag += 1
            if flag == 500:
                break


def val_resnet_test(net_param_path, file_path, res_name='./result.txt'):
    dataset_class = CaltechDataset(file_path, train=False)
    dataset = gen_dataset(dataset_class, train=False, batch_size=1)

    network = SPPLeNet5(num_class=256, num_channel=3, num_layers=5)

    # 加载已经保存的用于测试的模型
    param_dict = ms.load_checkpoint(net_param_path)
    # 加载参数到网络中
    ms.load_param_into_net(network, param_dict)

    model = ms.Model(network)

    res_pred = []
    for img_dict in dataset:
        imgs = img_dict['image']
        output = model.predict(ms.Tensor(imgs))
        predicted = np.argmax(output.asnumpy(), axis=1)
        print(predicted.item())
        res_pred.append(predicted.item()+1)
    res_file(res_pred, res_name=res_name)


if __name__ == '__main__':
    file_path = 'E:/Data/caltech_for_user'
    net_param_path = './resNet-adam.ckpt'
    val_resnet_train(net_param_path, file_path, train='train_val')

